import io
import logging
import numpy as np

from datetime import datetime, timedelta
from datetime import time as dttm
import matplotlib.pyplot as plt


# 设置中文字体
plt.rcParams["font.sans-serif"] = ["SimHei"]  # 用黑体显示中文
plt.rcParams["axes.unicode_minus"] = False  # 正常显示负号


def build_focus_windows(start_dt, end_dt):
    windows_list = []
    windows = []
    cur = datetime(start_dt.year, start_dt.month, start_dt.day)
    one_day = timedelta(days=1)
    while cur <= end_dt:
        for h0, h1 in [(9, 12), (12, 14), (14, 21)]:
            win_s = datetime(cur.year, cur.month, cur.day, h0, 0)
            win_e = datetime(cur.year, cur.month, cur.day, h1, 0)
            windows.append((win_s, win_e))
        cur += one_day
        windows.sort()
        windows_list.append(windows)
        windows = []
    return windows_list


def compress_dt(dt_obj, window, cum_offsets):
    for i, (s, e) in enumerate(window):
        if s <= dt_obj <= e:
            frac = (dt_obj - s).total_seconds() / (e - s).total_seconds()
            return cum_offsets[i] + frac * (cum_offsets[i + 1] - cum_offsets[i])
    return None


def get_free_time_slots(data):

    DARK_GREEN = "#2E7D32"
    BACKGROUND = "#C0C0C0"
    ALL_AVAILABLE_BLUE = "#0000FF"
    SHADOW_PATTERN = "////"  # 斜线阴影图案

    # 解析整体查询范围
    start_time = data["start_time"]
    end_time = data["end_time"]

    # 构建关注时段及压缩映射
    windows = build_focus_windows(start_time, end_time)
    if not windows:
        logging.info("查询时段非法")
        return ""
    cum_offsets = [0.0, 3.0, 3.5, 10.5]

    # 解析时间字符串函数
    def parse_time_slots(time_strings):
        slots_list = []
        slots = []
        if not time_strings:
            return slots
        a = time_strings[0][0].day
        for s in time_strings:
            if s[0].day == a:
                slots.append((s[0], s[1]))
            else:
                slots_list.append(slots)
                slots = []
                slots.append((s[0], s[1]))
                a = s[0].day
        slots_list.append(slots)
        return slots_list

    # 处理共同空闲和个人空闲
    common_raw = data.get("共同空闲时间")
    common_slots = parse_time_slots(common_raw) if isinstance(common_raw, list) else []
    person_slots = {}
    for p, ts in data.get("各人闲时间", {}).items():
        if p not in data.get("查询异常人员", []):
            if isinstance(ts, list):
                person_slots[p] = parse_time_slots(ts)
            else:
                person_slots[p] = []

    # 设置图形尺寸：宽度按人数
    persons = list(person_slots.keys())
    num_cols = len(windows)
    fig, ax = plt.subplots(figsize=(2 + num_cols * 1, 8.0))

    # 设置轴
    ax.set_xlim(0, num_cols)
    ax.set_ylim(10.5, 0)
    ax.xaxis.tick_top()
    ax.set_xticks(np.arange(num_cols) + 0.5)
    ax.set_xticklabels(
        [i[0][0].strftime("%m-%d") for i in windows], rotation=0, fontsize=10
    )
    x_labels = ax.xaxis.get_ticklabels()

    # Y 轴刻度与标签
    yticks, ylabels = [], []
    for s, e in windows[0]:
        cur = s
        if cur.hour not in (12, 21):
            while cur <= e:
                y = compress_dt(cur, windows[0], cum_offsets)
                if y is not None:
                    yticks.append(y)
                    ylabels.append(cur.strftime("%H:%M"))
                cur += timedelta(minutes=30)
    ax.set_yticks(yticks)
    ax.set_yticklabels(ylabels, fontsize=8)

    # 网格和列分隔线
    ax.grid(True, axis="y", linestyle="--", alpha=0.5)
    for x in range(1, num_cols):
        ax.axvline(x=x, color="gray", linestyle="-", linewidth=0.3, alpha=0.5)

    # 绘制条带函数
    def draw_slot(start, end, col_idx, color, hatch=None):
        cs = compress_dt(start, windows[col_idx], cum_offsets)
        ce = compress_dt(end, windows[col_idx], cum_offsets)
        if cs is not None and ce is not None:
            ax.add_patch(
                plt.Rectangle(
                    (col_idx, ce),
                    1,
                    cs - ce,
                    edgecolor=None,
                    facecolor=color,
                    alpha=0.9,
                    hatch=hatch,
                )
            )
        else:
            for ws, we in windows[col_idx]:
                ss = max(start, ws)
                ee = min(end, we)
                if ss < ee:
                    ys = compress_dt(ss, windows[col_idx], cum_offsets)
                    ye = compress_dt(ee, windows[col_idx], cum_offsets)
                    if ys is not None and ye is not None:
                        ax.add_patch(
                            plt.Rectangle(
                                (col_idx, ye),
                                1,
                                ys - ye,
                                edgecolor=None,
                                facecolor=color,
                                alpha=0.9,
                                hatch=hatch,
                            )
                        )

    no_work_slots = [(datetime(2025, 7, 8, 12, 0), datetime(2025, 7, 8, 14, 0))]
    if len(persons) > 1:
        # 绘制共同空闲
        for idx, common_slot in enumerate(common_slots):
            for st, et in common_slot:
                draw_slot(st, et, idx, DARK_GREEN)
            st, et = datetime(st.year, st.month, st.day, 12, 0), datetime(
                st.year, st.month, st.day, 14, 0
            )
            draw_slot(st, et, idx, BACKGROUND, hatch=SHADOW_PATTERN)
    else:
        # 绘制个人空闲
        for idx, _ in enumerate(person_slots.get(persons[0], [])):
            for st, et in person_slots.get(persons[0], [])[idx]:
                draw_slot(st, et, idx, DARK_GREEN)
            st, et = datetime(st.year, st.month, st.day, 12, 0), datetime(
                st.year, st.month, st.day, 14, 0
            )
            draw_slot(st, et, idx, BACKGROUND, hatch=SHADOW_PATTERN)

    # 底色（非关注背景）
    ax.add_patch(plt.Rectangle((0, 0), num_cols, 10.5, facecolor=BACKGROUND, alpha=0.4))

    # 异常人员说明
    exc_list = data.get("查询异常人员") or []
    if exc_list:
        exc_text = ", ".join(exc_list) if isinstance(exc_list, list) else str(exc_list)
        plt.figtext(
            0.5,
            0.01,
            f"异常人员: {exc_text}",
            ha="center",
            fontsize=10,
            bbox={"facecolor": "#FF9800", "alpha": 0.7, "pad": 4},
        )

    title_range = f"{start_time.strftime('%m-%d')} to {end_time.strftime('%m-%d')}"
    plt.title(
        f"{'公共空闲' if len(persons) > 1 else '['+persons[0]+']空闲'} ({title_range})",
        pad=10,
        fontsize=12,
    )
    plt.tight_layout(rect=[0, 0.05, 1, 0.95])

    # 图例
    from matplotlib.patches import Patch

    legend_elems = [
        Patch(facecolor=DARK_GREEN, alpha=0.8, label="空闲时间"),
        Patch(facecolor=BACKGROUND, alpha=0.8, label="非空闲时间"),
    ]
    ax.legend(
        handles=legend_elems,
        loc="upper right",
        bbox_to_anchor=(1, 1 + (1.2 / fig.get_figheight())),
        fontsize=8,
        ncol=3,
    )

    # 输出 PNG
    buf = io.BytesIO()
    fig.canvas.print_png(buf)
    img_data = buf.getvalue()
    buf.close()
    return img_data


if __name__ == "__main__":
    # 输入数据
    data = {
        "各人闲时间": {
            "dianlv": [
                (datetime(2025, 7, 8, 14, 0), datetime(2025, 7, 8, 15, 0)),
                (datetime(2025, 7, 8, 16, 0), datetime(2025, 7, 8, 18, 0)),
                (datetime(2025, 7, 9, 14, 30), datetime(2025, 7, 9, 15, 30)),
                (datetime(2025, 7, 9, 16, 0), datetime(2025, 7, 9, 21, 0)),
            ],
        },
        "共同空闲时间": [
            (datetime(2025, 7, 8, 14, 0), datetime(2025, 7, 8, 15, 0)),
            (datetime(2025, 7, 8, 16, 0), datetime(2025, 7, 8, 18, 0)),
            (datetime(2025, 7, 9, 14, 30), datetime(2025, 7, 9, 15, 30)),
            (datetime(2025, 7, 9, 16, 0), datetime(2025, 7, 9, 21, 0)),
        ],
        "查询异常人员": ["jianweiwang"],
        "start_time": datetime(2025, 7, 8, 10, 0),
        "end_time": datetime(2025, 7, 9, 12, 0),
    }
    with open("hhh.png", mode="wb") as f:
        f.write(get_free_time_slots(data))
